#include "Hook.hpp"
#include "../base/Defind.hpp"
#include "../base/BaseLog.hpp"
#include "FdManager.hpp"
#include "Fiber.hpp"
#include "IOManager.hpp"
#include "TimerManager.hpp"

#include <asm-generic/ioctls.h>
#include <asm-generic/socket.h>
#include <bits/types/struct_timeval.h>
#include <cerrno>
#include <cstdarg>
#include <cstdint>
#include <ctime>
#include <dlfcn.h>
#include <fcntl.h>
#include <memory>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <utility>

namespace DBServer::coroutine {

    void hook_init();

    // TCP连接超时时间
    static std::uint64_t tcp_connect_timeout = TCP_CONNECT_TIMEOUT;

    static std::uint64_t s_connect_timeout = -1;

    // 提供线程粒度的hook函数
    static thread_local bool hook_enable = []() {
        hook_init();
        return false;
    }();

#define HOOK_FUN(XX)     \
    XX(sleep)            \
    XX(usleep)           \
    XX(nanosleep)        \
    XX(socket)           \
    XX(connect)          \
    XX(accept)           \
    XX(read)             \
    XX(recv)             \
    XX(readv)            \
    XX(recvfrom)         \
    XX(recvmsg)          \
    XX(write)            \
    XX(writev)           \
    XX(send)             \
    XX(sendto)           \
    XX(sendmsg)          \
    XX(close)            \
    XX(fcntl)            \
    XX(ioctl)            \
    XX(getsockopt)       \
    XX(setsockopt)

    void hook_init() {
        static bool is_inited = false;
        if (is_inited) {
            return;
        }

#define XX(name) name##_f = (name##_fun)dlsym(RTLD_NEXT, #name);
        HOOK_FUN(XX)
#undef XX
        is_inited = true;
    }

    bool is_hook_enable() {
        return hook_enable;
    }

    void set_hook_enable(bool flag) {
        hook_enable = flag;
    }
}

struct timer_info {
    int cancelled = 0;
};

template<typename OriginFun, typename ...Args>
ssize_t do_io(int fd, OriginFun fun, const char *hook_fu_name, std::uint32_t event, int timeout_so, Args&&... args) {

    // 该线程没有打开hook
    if (!DBServer::coroutine::hook_enable) {
        return fun(fd, std::forward<Args>(args)...);
    }

    // fd没有被注册在fd manager中
    auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(fd);
    if (!ctx) {
        return fun(fd, std::forward<Args>(args)...);
    }

    // fd已经被关闭
    if (ctx->attr_close()) {
        errno = EBADF;
        return -1;
    }

    // fd 不是socket 或者fd是由用户自己管理的
    if (!ctx->attr_socket() || ctx->attr_user_nonblock()) {
        return fun(fd, std::forward<Args>(args)...);
    }


    std::uint64_t to = ctx->get_timeout(timeout_so);
    auto tinfo = std::make_shared<timer_info>();

retry:
    ssize_t n;
    do
    {
        n = fun(fd, std::forward<Args>(args)...);
    } while (n == -1 && errno == EINTR);

    if (n == -1 && errno == EAGAIN) 
    {
        auto iom = DBServer::coroutine::IOManager::GetThis();
        DBServer::coroutine::Timer::ptr timer;
        std::weak_ptr<timer_info> winfo(tinfo);

        if (to != (std::uint64_t)-1) {
            timer = iom->add_condition_timer(to, [winfo, fd, iom, event](){
                auto t = winfo.lock();
                if (!t || t->cancelled) {
                    return;
                }
                t->cancelled = ETIMEDOUT;
                // 取消注册的事件
                iom->cancel_ev(fd, (DBServer::coroutine::IOEvent)event);
            }, winfo);
        }

        int rt = iom->add_ev(fd, (DBServer::coroutine::IOEvent)event);
        if (rt) {
            error("%s add_event(%d)\n", hook_fu_name, fd);
            if (timer) {
                timer->cancel();
            }
            return -1;
        }
        else
        {
            DBServer::coroutine::Fiber::Yield();
            if (timer) {
                timer->cancel();
            }
            if (tinfo->cancelled) {
                errno = tinfo->cancelled;
                return -1;
            }
            goto retry;
        }
    }
    return n;
}

extern "C" {
#define XX(name) name##_fun name##_f = nullptr;
    HOOK_FUN(XX);
#undef XX

unsigned int sleep(unsigned int seconds) {
    if (!DBServer::coroutine::hook_enable) {
        return sleep_f(seconds);
    }

    auto fiber = DBServer::coroutine::Fiber::GetThis();
    auto iom = DBServer::coroutine::IOManager::GetThis();
    iom->add_timer(seconds * 1000, [iom, fiber]() {
        iom->schedule(fiber, -1);
    });
    DBServer::coroutine::Fiber::Yield();
    return 0;
}

int usleep(unsigned int usec) {
    if (DBServer::coroutine::hook_enable) {
        return usleep_f(usec);
    }
    auto fiber = DBServer::coroutine::Fiber::GetThis();
    auto iom = DBServer::coroutine::IOManager::GetThis();
    iom->add_timer(usec / 1000, [iom, fiber](){
        iom->schedule(fiber, -1);
    });
    return 0;
}

int nanosleep(const struct timespec *req, struct timespec *rem) {
    if (!DBServer::coroutine::hook_enable) {
        return nanosleep_f(req, rem);
    }

    int timeout_ms = req->tv_sec * 1000 + req->tv_nsec / 1000 / 1000;
    auto fiber = DBServer::coroutine::Fiber::GetThis();
    auto iom = DBServer::coroutine::IOManager::GetThis();
    iom->add_timer(timeout_ms, [iom, fiber](){
        iom->schedule(fiber, -1);
    });
    DBServer::coroutine::Fiber::Yield();
    return 0;
}

int socket(int domain, int type, int protocol) {
    if (!DBServer::coroutine::hook_enable) {
        return socket_f(domain, type, protocol);
    }

    int fd = socket_f(domain, type, protocol);
    if (fd == -1) {
        return fd;
    }
    DBServer::coroutine::FdManager::GetInstance()->reg(fd);
    return fd;
}

int connect_with_timeout(int fd, const struct sockaddr *addr, socklen_t addrlen, uint64_t timeout_ms) {
    if (DBServer::coroutine::hook_enable) {
        return connect_f(fd, addr, addrlen);
    }

    auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(fd);
    if (!ctx || ctx->attr_close()) {
        errno = EBADF;
        return -1;
    }

    if (!ctx->attr_socket()) {
        return connect_f(fd, addr, addrlen);
    }

    if (ctx->attr_user_nonblock()) {
        return connect_f(fd, addr, addrlen);
    }

    int n = connect_f(fd, addr, addrlen);
    if (n == 0) {
        return 0;
    } else if (n != -1 || errno != EINPROGRESS) {
        return n;
    }

    auto iom = DBServer::coroutine::IOManager::GetThis();
    DBServer::coroutine::Timer::ptr timer;
    auto tinfo = std::make_shared<timer_info>();
    std::weak_ptr<timer_info> winfo(tinfo);

    if (timeout_ms != (uint64_t)-1) {
        timer = iom->add_condition_timer(timeout_ms, [winfo, fd, iom](){
            auto t = winfo.lock();
            if (!t || t->cancelled) {
                return;
            }
            t->cancelled = ETIMEDOUT;
            iom->cancel_ev(fd, DBServer::coroutine::IOEvent::WRITE);
        }, winfo);
    }

    int rt = iom->add_ev(fd, DBServer::coroutine::IOEvent::WRITE);
    if (rt == 0) {
        DBServer::coroutine::Fiber::Yield();
        if (timer) {
            timer->cancel();
        } 
        if (tinfo->cancelled) {
            errno = tinfo->cancelled;
            return -1;
        }
    } else {
        if (timer) {
            timer->cancel();
        }
        error("connect add_ev(%d, WRITE) error", fd);
    }

    int error = 0;
    socklen_t len = sizeof(int);
    if (-1 == getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len)) {
        return -1;
    }
    if (!error) {
        return 0;
    } else {
        errno = error;
        return -1;
    }
}

int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
    return connect_with_timeout(sockfd, addr, addrlen, DBServer::coroutine::s_connect_timeout);
}

using IOEvent = DBServer::coroutine::IOEvent;

int accept(int s, struct sockaddr *addr, socklen_t *addrlen) {
    int fd = do_io(s, accept_f, "accept", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, addr, addrlen);
    if (fd >= 0) {
        DBServer::coroutine::FdManager::GetInstance()->reg(fd);
    }
    return fd;
}

ssize_t read(int fd, void *buf, size_t count) {
    return do_io(fd, read_f, "read", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, buf, count);
}

ssize_t readv(int fd, const struct iovec *iov, int iovcnt) {
    return do_io(fd, readv_f, "readv", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, iov, iovcnt);
}

ssize_t recv(int sockfd, void *buf, size_t len, int flags) {
    return do_io(sockfd, recv_f, "recv", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, buf, len, flags);
}

ssize_t recvfrom(int sockfd, void *buf, size_t len, int flags, struct sockaddr *src_addr, socklen_t *addrlen) {
    return do_io(sockfd, recvfrom_f, "recvfrom", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, buf, len, flags, src_addr, addrlen);
}

ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags) {
    return do_io(sockfd, recvmsg_f, "recvmsg", (std::uint32_t)IOEvent::READ, SO_RCVTIMEO, msg, flags);
}

ssize_t write(int fd, const void *buf, size_t count) {
    return do_io(fd, write_f, "write", (std::uint32_t)IOEvent::WRITE, SO_SNDTIMEO, buf, count);
}

ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
    return do_io(fd, writev_f, "writev", (std::uint32_t)IOEvent::WRITE, SO_SNDTIMEO, iov, iovcnt);
}

ssize_t send(int s, const void *msg, size_t len, int flags) {
    return do_io(s, send_f, "send", (std::uint64_t)IOEvent::WRITE, SO_SNDTIMEO, msg, len, flags);
}

ssize_t sendto(int s, const void *msg, size_t len, int flags, const struct sockaddr *to, socklen_t tolen) {
    return do_io(s, sendto_f, "sendto", (std::uint32_t)IOEvent::WRITE, SO_SNDTIMEO, msg, len, flags, to, tolen);
}

ssize_t sendmsg(int s, const struct msghdr *msg, int flags) {
    return do_io(s, sendmsg_f, "sendmsg", (std::uint32_t)IOEvent::WRITE, SO_SNDTIMEO, msg, flags);
}

int close(int fd) {
    if (!DBServer::coroutine::hook_enable) {
        return close_f(fd);
    }

    auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(fd);
    if (ctx) {
        auto iom = DBServer::coroutine::IOManager::GetThis();
        if (iom) {
            iom->cancel_all(fd);
        }
        DBServer::coroutine::FdManager::GetInstance()->del(fd);
    }
    return close_f(fd);
}

int fcntl(int fd, int cmd, ...) {
    va_list va;
    va_start(va, cmd);
    switch (cmd) {
        case F_SETFL: 
            {
                int arg = va_arg(va, int);
                va_end(va);
                auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(fd);
                if (!ctx || ctx->attr_close() || ctx->attr_socket()) {
                    return fcntl_f(fd, cmd, arg);
                }
                ctx->set_user_nonblock(arg & O_NONBLOCK);
                if (ctx->attr_sys_nonblock()) {
                    arg |= O_NONBLOCK;
                } else {
                    arg &= ~O_NONBLOCK;
                }
                return fcntl_f(fd, cmd, arg);
            }
            break;
        case F_GETFL:
            {
                va_end(va);
                int arg = fcntl_f(fd, cmd);
                auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(fd);
                if (!ctx || ctx->attr_close() || !ctx->attr_socket()) {
                    return arg;
                }
                if (ctx->attr_user_nonblock()) {
                    return arg | O_NONBLOCK;
                } else {
                    return arg & ~O_NONBLOCK;
                }
            }
            break;
        case F_DUPFD:
        case F_DUPFD_CLOEXEC:
        case F_SETFD:
        case F_SETOWN:
        case F_SETSIG:
        case F_SETLEASE:
        case F_NOTIFY:
#ifdef F_SETPIPE_SZ
        case F_SETPIPE_SZ:
#endif
            {
                int arg = va_arg(va, int);
                va_end(va);
                return fcntl_f(fd, cmd, arg);
            }
            break;
        case F_SETLK:
        case F_SETLKW:
        case F_GETLK:
            {
                struct flock *arg = va_arg(va, struct flock *);
                va_end(va);
                return fcntl_f(fd, cmd, arg);
            }
        case F_GETOWN_EX:
        case F_SETOWN_EX:
            {
                struct f_owner_exlock *arg = va_arg(va, struct f_owner_exlock *);
                va_end(va);
                return fcntl_f(fd, cmd, arg);
            }
            break;
        default:
            va_end(va);
            return fcntl_f(fd, cmd);
        }
}

int ioctl(int d, unsigned long int request, ...) {
    va_list va;
    va_start(va, request);
    void *arg = va_arg(va, void*);
    va_end(va);

    if (FIONBIO == request) {
        bool user_nonblock = !!*(int *)arg;
        auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(d);
        if (!ctx || ctx->attr_close() || !ctx->attr_socket()) {
            return ioctl_f(d, request, arg);
        }
        ctx->set_user_nonblock(user_nonblock);
    }
    return ioctl_f(d, request, arg);
}

int getsockopt(int sockfd, int level, int optname, void *optval, socklen_t *optlen) {
    return getsockopt_f(sockfd, level, optname, optval, optlen);
}

int setsockopt(int sockfd, int level, int optname, const void *optval, socklen_t optlen) {
    if (DBServer::coroutine::hook_enable) {
        return setsockopt_f(sockfd, level, optname, optval, optlen);
    }
    if (level == SOL_SOCKET) {
        if (optname == SO_RCVTIMEO || optname == SO_SNDTIMEO) {
            auto ctx = DBServer::coroutine::FdManager::GetInstance()->get(sockfd);
            if (ctx) {
                const timeval *v = (const timeval*)optval;
                ctx->set_timeout(optname, v->tv_sec * 1000 + v->tv_usec / 1000);
            }
        }
    }
    return setsockopt_f(sockfd, level, optname, optval, optlen);
}
}